#include "imagebind_inference.h"
#include <algorithm>
#include <string>
#include <numeric>
#include <json/json.h>
#include <fstream>
#include <dirent.h>
#include <iostream>
#include <chrono>
#include "cuda_runtime_api.h"
#include "logging.h"
#include <opencv2/opencv.hpp>
#include "NvInfer.h"
#include <fstream>

static Logger gLogger;

std::map<std::string ,std::vector<float>> load_pre_text_file(std::vector<std::string> &files) {
      // std::cout<<"&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&"<<std::endl;
      std::map<std::string ,std::vector<float>> known_embedding;
      for (auto &file : files)
      {
          size_t pos1 = file.find_last_of('/');
          // std::cout<<"pos1:"<<pos1<<std::endl;
          std::string folder = file.substr(0, pos1);
          // std::cout<<"folder:"<<folder<<std::endl;
          std::string cls = folder.substr(folder.find_last_of('/')+1, folder.length());
          // std::cout<<"cls:"<<cls<<std::endl;
          size_t pos2 = file.find('.');
          std::string index = file.substr(pos1+1, pos2-pos1-1);
          // std::cout<<"index:"<<index<<std::endl;
          std::ifstream in_file(file);
          std::vector <float> record;
          while (in_file)
          {
              std::string s;
              if (!getline(in_file, s)) break;
              std::istringstream ss(s);
              while (ss)
              {
                  std::string s;
                  if (!getline(ss, s, ',')) break;
                  record.push_back(atof(s.c_str()));
              }
          }

          known_embedding[cls+"_"+index].assign(record.begin(), record.end());
      }
      return known_embedding;
}

bool load_pre_text_embeddings(std::string path, std::vector<std::map<std::string ,std::vector<float>>> &known_embeddings, std::vector<std::string> cls_names) {
    //file handle
    DIR *d = opendir(path.c_str());
    //file path
    std::vector<std::string> files;
    std::string p;
    if (d != NULL)
    {
      struct dirent *dt = NULL;
      while (dt = readdir(d))
      {
        if (strcmp(dt->d_name, ".") != 0 && strcmp(dt->d_name, "..") != 0)
        {
          if (dt->d_type == 4 )
          {
            if (std::find(cls_names.begin(), cls_names.end(), dt->d_name) != cls_names.end())
            {
              std::cout<<"find pretrain vector class = "<<dt->d_name<<std::endl;
              load_pre_text_embeddings(p.assign(path).append("/").append(dt->d_name), known_embeddings, cls_names);
            }
          } else
          {
            files.push_back(p.assign(path).append("/").append(dt->d_name));
          }
        }
      }

      for (auto lfilename: files){
        std::cout<<"lfilename:"<<lfilename<<std::endl;
      }

      closedir(d);
      if(!files.empty())
      {
        known_embeddings.push_back(load_pre_text_file(files));
      }
    }
    return true;
}

std::vector<float> softmax(std::vector<float> &conf, float temperature)
  {
    for (auto &value : conf)
    {
      value *= temperature;
    }

    std::vector<float>::iterator biggest = std::max_element(conf.begin(),conf.end());
    float max_value = *biggest;
    std::cout<<"maxval:"<<max_value<<std::endl;
    
    float total = 0;
    for (auto x : conf)
    {
      total += exp(x-max_value);
    }
    std::cout<<"total:"<<total<<std::endl;

    std::vector<float> result;
    for (auto x : conf)
    {
       std::cout<<"newval:"<<exp(x-max_value)<<std::endl;
      result.push_back(exp(x-max_value)/total);
    }
    return result;
}

imagebindInfer::imagebindInfer(){
    std::cout<<"init empty imagebindInfer"<<std::endl;
}

imagebindInfer::~imagebindInfer(){
    context->destroy();
    engine->destroy();
    runtime->destroy();
}

int read_files_in_dir(const char *p_dir_name, std::vector<std::string> &file_names) {
    DIR *p_dir = opendir(p_dir_name);
    if (p_dir == nullptr) {
        return -1;
    }

    struct dirent* p_file = nullptr;
    while ((p_file = readdir(p_dir)) != nullptr) {
        if (strcmp(p_file->d_name, ".") != 0 &&
                strcmp(p_file->d_name, "..") != 0) {
            //std::string cur_file_name(p_dir_name);
            //cur_file_name += "/";
            //cur_file_name += p_file->d_name;
            std::string cur_file_name(p_file->d_name);
            file_names.push_back(cur_file_name);
        }
    }

    closedir(p_dir);
    return 0;
}

bool cv2vec(const cv::Mat src_image, std::vector<float> &input_vectors) {
      cv::Mat channels[3]; //借用来进行HWC->CHW
      cv::split(src_image, channels);
      for (int i = 0; i < src_image.channels(); i++)  // HWC->CHW
      {
          std::vector<float> data = std::vector<float>(channels[i].reshape(1, src_image.cols * src_image.rows));
          input_vectors.insert(input_vectors.end(), data.begin(), data.end());
      }
      
      return true;
  }

bool image_preprocess(const cv::Mat &src_image, int width, int height, std::vector<float> &input_vectors, std::vector<float>&means, std::vector<float>&stds) {
    cv::Mat src_temp, dst_temp, dst_image;
    // 1. resize by CUBIC mode
    cv::resize(src_image, src_temp, cv::Size(width, height), 0, 0, cv::INTER_CUBIC);
    src_temp.convertTo(dst_temp, CV_32F);
    dst_temp.copyTo(dst_image);

    // 2. convert BGR to RGB
    cv::cvtColor(dst_image, dst_image, cv::COLOR_BGR2RGB);
    for (int i = 0; i < dst_image.rows; i++) {
        for (int j = 0; j < dst_image.cols; j++) {
            cv::Vec3f buf = dst_image.at<cv::Vec3f>(i,j);
            // printf("buf %f %f %f\n", buf[0], buf[1], buf[2]);
            buf[0] = ((buf[0] / 255.0f) - means[0]) / stds[0];
            buf[1] = ((buf[1] / 255.0f) - means[1]) / stds[1];
            buf[2] = ((buf[2] / 255.0f) - means[2]) / stds[2];
            dst_image.at<cv::Vec3f>(i,j) = buf;
            // printf("buf %f %f %f\n", buf[0], buf[1], buf[2]);
        }
    }

    // 4.convert mat to vector
    cv2vec(dst_image, input_vectors);

    return true;
  }



bool imagebindInfer::initModelConfig(const std::string modelconfigpath){
    Json::Reader json_reader;
    Json::Value rootValue;
    std::ifstream infile(modelconfigpath.c_str(), std::ios::binary);
    if (!infile.is_open()){
        std::cout<<"init imagebind config file is open false"<<std::endl;
        return false;
    }

    if (json_reader.parse(infile, rootValue)){
        _clip_Width = rootValue["width"].asInt();
        _clip_Height =  rootValue["height"].asInt();
        text_emdeding_path = rootValue["text_embeding"].asString();
        Json::Value targetVals = rootValue["class_names"];
        for(unsigned int i=0; i < targetVals.size(); i++){
            imagebind_class_names.insert(targetVals[i].asString());
        }
    
    }else{
        std::cout<<"Error: Can not parse imagebind config file"<<std::endl;
        return false;
    }
    return true;

}


bool imagebindInfer::initModelRunCon(const std::string modelRunconfigpath){
    Json::Reader json_reader;
    Json::Value rootValue;

    // 6. load pre-save embeddings
    std::string file_path = "/opt/micros/share/capability_package/target_name_map.json";
    std::cout<<"json file path = "<<file_path<<std::endl;
    std::ifstream jsonfile(file_path, std::ios::binary);
    if (!jsonfile.is_open())
    {
        std::cout<<"Error open json file = "<<file_path<<std::endl;
        return false;
    }
    
    if((json_reader.parse(jsonfile, rootValue)))
    {
        Json::Value::Members member = rootValue.getMemberNames();
        for (Json::Value::Members::iterator iter = member.begin(); iter != member.end(); ++iter)
        {
            cls_names.push_back(rootValue[*iter].asString());
        }
    }else{
        std::cout<<"Error: Can not parse imagebind run config file"<<std::endl;
        return false;
    }
    jsonfile.close();

    std::ifstream infile(modelRunconfigpath.c_str(), std::ios::binary);
    if (!infile.is_open()){
        std::cout<<"init imagebind run config file is open false"<<std::endl;
        return false;
    }

    if (json_reader.parse(infile, rootValue)){
        Json::Value meanVals = rootValue["mean"];
        Json::Value stdVals = rootValue["std"];
        for(unsigned int i=0; i < meanVals.size(); i++){
            means.push_back(meanVals[i].asFloat());
            stds.push_back(stdVals[i].asFloat());
        }

        Json::Value targetVals = rootValue["target_names"];
        for(unsigned int i=0; i < targetVals.size(); i++){
            target_names.insert(targetVals[i].asString());
        } 

        Json::Value textVals = rootValue["text_embeding_name"];
        for(unsigned int i=0; i < textVals.size(); i++){
            cls_names.push_back(textVals[i].asString());
        } 
        temperature = rootValue["temperature"].asFloat();
        std::cout<<"temperature0:"<<temperature<<std::endl;
        pattern = rootValue["pattern"].asInt();
    }else{
        std::cout<<"Error: Can not parse imagebind run config file"<<std::endl;
        return false;
    }
    infile.close();

    for (auto jname: cls_names)
    {
        std::cout<<"jsonname:"<<jname<<std::endl;
    }
    return true;
}

bool imagebindInfer::initEngine(const std::string enginePath){
    char *trtModelStream{nullptr};
    size_t size{0};
    std::ifstream file(enginePath);
    if (file.good()) {
        file.seekg(0, file.end);
        size = file.tellg();
        file.seekg(0, file.beg);
        trtModelStream = new char[size];
        assert(trtModelStream);
        file.read(trtModelStream, size);
        file.close();
    }

    IRuntime* runtime = createInferRuntime(gLogger);
    assert(runtime != nullptr);
    ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size);
    assert(engine != nullptr);
    IExecutionContext* context = engine->createExecutionContext();
    nvinfer1::Dims input_dims = nvinfer1::Dims4(1, 3, _clip_Width, _clip_Height);
    context->setBindingDimensions(0, input_dims);
    nvinfer1::Dims out_dims = nvinfer1::Dims2(1, 6);
    context->setBindingDimensions(1, out_dims);
    assert(context != nullptr);
    delete[] trtModelStream;
}


bool imagebindInfer::Initialize(const std::string& configPath){
    bool flag;
    std::string model_path =  configPath + "/cliptiny.engine";
    std::string imagebind_config = configPath + "/imagebind_config.json";
    std::string imagebindRun_config = configPath + "/imagebind_run.json";

    flag = initModelConfig(imagebind_config);
    if (!flag){
      std::cout<<"init Model false"<<std::endl;
      return false;
    }
    flag = initModelRunCon(imagebindRun_config);
    if (!flag){
      std::cout<<"init model Run config false"<<std::endl;
      return false;
    }
    flag = initEngine(model_path);
    if (!flag){
      std::cout<<"init Engine false"<<std::endl;
      return false;
    }
    
    load_pre_text_embeddings(text_emdeding_path, _known_embeddings, cls_names);
}


bool imagebindInfer::inference(cv::Mat &visible_pFrame,yoloV3GPUDetector &yolo,std::vector<cv::Rect2d> &pRoi,std::vector<std::string> &pName,std::vector<float> &pConf){
    if (!first_init){
      yolo.setParams(0.2, target_names,"",1024,768);
      first_init = true;
    }
    if (visible_pFrame.empty())
    {
      std::cout<<"The image is empty!!!!!"<<std::endl;
      return false;
    }
    
    bool detected = yolo.inference(visible_pFrame, pRoi, pName, pConf);
    if (!detected)
    {
      return false;
    }
    std::cout<<"pName = " << pName.size()<<" pConf = "<< pConf.size()<<std::endl;
    for (int i=0; i<pName.size(); i++)
    {
      std::cout<<"i = "<<i<<" pName = "<<pName[i]<<" pConf = "<<pConf[i]<<std::endl;
    }
    // maximum processed target number
    int max_target_num = std::min(2, int(pRoi.size()));
    std::cout<<"max_target_num = "<< max_target_num<<std::endl;
    int modify_count = 0;
    std::cout<<"size = "<< _known_embeddings.size()<<std::endl;
    std::vector<int> rm_index;
    if (_known_embeddings.size() > 1)
    {
      for (int i = 0; i < max_target_num; i++)
      {
        int index;
        if(imagebind_detect_order == 1)
        {
          index = pRoi.size() - i - 1;
          std::cout<<"ascend order: -------------------------------------------------index = "<< index<<std::endl;
        } else 
        {
          index = i;
          std::cout<<"descend order: -------------------------------------------------index = "<< index<<std::endl;
        }
        cv::Rect2d Roi = pRoi[index];
        std::string Name = pName[index];
        float Conf = pConf[index];
        std::cout<<"Yolo detect category: "<< Name.c_str()<<" scores: "<< Conf<<std::endl;

        std::unordered_set<std::string>::iterator it = imagebind_class_names.find(Name);
        if(it ==imagebind_class_names.end()){
          std::cout<<"imagebind does not support predicting: "<< Name.c_str()<<std::endl;
          continue;
        }
        
        int x1 = std::max(0, int(Roi.x - 10));
        int y1 = std::max(0, int(Roi.y - 10));
        int pwidth = std::min(int(x1 + Roi.width + 20), int(visible_pFrame.cols -1)) - x1;
        int pheight = std::min(int(y1 + Roi.height + 20), int(visible_pFrame.rows -1)) - y1;

        cv::Rect largeRect(x1, y1, pwidth, pheight);
        cv::Mat image_roi = visible_pFrame(largeRect);
        std::vector<float> input_vectors;
        for (auto val: means){
          std::cout<<"mean:"<<val<<std::endl;
        }

        image_preprocess(image_roi, _clip_Width, _clip_Height, input_vectors, means, stds);

        std::vector<float> _embeddings(6); 
        const ICudaEngine& engine = context->getEngine();
        assert(engine.getNbBindings() == 2);
        void* buffers[2];
        const int inputIndex = engine.getBindingIndex("input");
        const int outputIndex = engine.getBindingIndex("output");
        std::cout<<"inputIndex:"<<inputIndex<<", outputIndex"<<outputIndex<<std::endl;
        // Create GPU buffers on device
        CHECK(cudaMalloc(&buffers[inputIndex],  3 * _clip_Width * _clip_Height * sizeof(float)));
        CHECK(cudaMalloc(&buffers[outputIndex], 6 * sizeof(float)));

        // Create stream
        cudaStream_t stream;
        CHECK(cudaStreamCreate(&stream));

        // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
        CHECK(cudaMemcpyAsync(buffers[inputIndex], input_vectors.data(),  3 * _clip_Width * _clip_Height * sizeof(float), cudaMemcpyHostToDevice, stream));
        context->enqueue(1, buffers, stream, nullptr);
        CHECK(cudaMemcpyAsync(_embeddings.data(), buffers[outputIndex],  6 * sizeof(float), cudaMemcpyDeviceToHost, stream));
        cudaStreamSynchronize(stream);

        // Release stream and buffers
        cudaStreamDestroy(stream);
        CHECK(cudaFree(buffers[inputIndex]));
        CHECK(cudaFree(buffers[outputIndex]));

        printf("IMAGEBIND: output = %f %f %f %f %f %f\n", _embeddings[0], _embeddings[1], _embeddings[2], _embeddings[3], _embeddings[4], _embeddings[5]);
        std::vector<float> conf;
        std::vector<std::string> name;
        int count = 0;
        int background_index = 0;
        int others_index = 0;
        for (auto &known_embedding:_known_embeddings)
        {
          float cls_in_conf=-1000;
          std::string cls_in_name;
          for (auto &it : known_embedding)
          {
            float val = std::inner_product(_embeddings.begin(), _embeddings.end(), it.second.begin(), 0.0);
            if (val > cls_in_conf)
            {
              cls_in_conf = val;
            }
            cls_in_name = it.first.substr(0, it.first.find_last_of("_"));
          }
          conf.push_back(cls_in_conf);
          name.push_back(cls_in_name);
          if (cls_in_name == "Background"){
              background_index=count;
          }
          if (cls_in_name == "Others"){
              others_index=count;
          }
          count++;
        }

        for(auto val: name){
          std::cout<<"name:"<<val<<std::endl;
        }

        std::cout<<"background_index:"<<background_index<<",          others_index:"<<others_index<<" temperature:"<<temperature<<std::endl;
        std::vector<float> softmax_conf = softmax(conf, temperature);
        // cal max value and index
        std::vector<float>::iterator biggest = std::max_element(softmax_conf.begin(), softmax_conf.end());
        float max_value = *biggest;
        int max_value_index = std::distance(softmax_conf.begin(), biggest);

        // std
        auto size = std::distance(softmax_conf.begin(), softmax_conf.end());
        double avg = std::accumulate(softmax_conf.begin(), softmax_conf.end(),0.0) / size;
        double variance(0);
        std::for_each(softmax_conf.begin(), softmax_conf.end(), [avg, &variance](const float &num){variance += (num - avg) * (num -avg);});
        variance /= size;
        auto standard = std::sqrt(variance);

        std::cout<<"IMAGEBIND: Current detection box category: "<<Name<<", ImageBind category: "<<name[max_value_index]<<std::endl;
        for (int j = 0; j < softmax_conf.size(); j++)
        {
          std::cout<<"IMAGEBIND: index = "<<j<<" name: "<<name[j]<<" softmax_conf: "<<softmax_conf[j]<<std::endl;
        }

        if (pattern == 1)
        {
          std::cout<<"pattern:"<<pattern<<std::endl;
          if (max_value_index == background_index || max_value_index == others_index)// || standard < 0.2)
          { 
            modify_count++;
            rm_index.push_back(index);
            std::cout<<"TARGET REMOVED!!! Current detection box category: "<<Name<<", ImageBind Max category confidence is:"<<max_value<<", this box should be False Positive (background)"<<std::endl;
          }else
          {
              modify_count++;
              pName[index] = name[max_value_index];
              pConf[index] = max_value;
              std::cout<<"TARGET MODIFIED!!! Current detection box category: "<<Name<<", ImageBind category: "<<name[max_value_index]<<" ImageBind confidence is "<< max_value<<std::endl;
          }
        } else{
            std::string imgbind_name = name[max_value_index];
            if ((max_value < 0.2) || max_value_index == background_index || max_value_index == others_index){
              modify_count++;
              rm_index.push_back(index);
              std::cout<<"TARGET REMOVED!!! Current detection box category: "<<Name<<", ImageBind Max category confidence is:"<<max_value<<", this box should be False Positive (background)"<<std::endl;

            }else if(((Name=="Hummer") || (Name=="Himars_launcher") || (Name=="Command") || (Name=="Himars_carrier") || (Name=="Cross") || (Name=="Others")) && (imgbind_name != Name) && (max_value > Conf)){
                modify_count++;
                pName[index] = name[max_value_index];
                pConf[index] = max_value;
                std::cout<<"TARGET MODIFIED!!! Current detection box category: "<<Name<<", ImageBind category: "<<name[max_value_index]<<" ImageBind confidence is "<< max_value<<std::endl;
            }else {
              std::cout<<"TARGET KEPT!!!Current detection box category: "<<Name<<" ImageBind category: " <<name[max_value_index]<<std::endl;
            }
           
        }
      }
    } else 
    {
      std::cout<<"DO NOT have enough pretrained text embeddings, ImageBind no inference"<<std::endl;
    }

    if (rm_index.size() == pRoi.size())
    {
      if (imagebind_detect_order == 0)
      {
        rm_index.erase(rm_index.begin());
      } else
      {
        rm_index.erase(rm_index.begin()+rm_index.size()-1);
      }
    }
    
    // remove error Roi
    if (rm_index.size() > 0)
    {
      for (int i = 0; i < rm_index.size(); i++)
      {
        if (imagebind_detect_order == 0)
        {
          pRoi.erase(pRoi.begin()+rm_index[i]-i);
          pName.erase(pName.begin()+rm_index[i]-i);
          pConf.erase(pConf.begin()+rm_index[i]-i);
        } else
        {
          pRoi.erase(pRoi.begin()+rm_index[i]);
          pName.erase(pName.begin()+rm_index[i]);
          pConf.erase(pConf.begin()+rm_index[i]);
        }
      }
    }

    std::cout<<"IMAGEBIND: ImageBind modified boxes numbers: "<< modify_count<<std::endl;
    if (pRoi.size() == 0 && detected == true)
    {
      detected = false;
    }
    
    return detected;

}

